package org.dllwh.milvus.v1;

import cn.hutool.core.collection.CollectionUtil;
import com.alibaba.fastjson2.JSONObject;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.MutationResult;
import io.milvus.param.R;
import io.milvus.param.dml.DeleteParam;
import io.milvus.param.dml.InsertParam;
import lombok.extern.slf4j.Slf4j;

import java.util.List;

/**
 * 把今天最好的表现当作明天最新的起点．．～
 * <p>
 * Today the best performance as tomorrow the newest starter!
 *
 * @ClassName MilvusDataOperateHelper
 * @Description Milvus V1 数据操作
 * @Author: <a href="mailto:duleilewuhen@sina.com">独泪了无痕</a>
 * @Date 2025-05-24 14:35
 * @版本 V 1.0.1
 * @since JDK 17
 */
@Slf4j
public class MilvusDataOperateHelper {
    /**
     * 插入数据
     *
     * @param collectName 指定集合名称
     * @param fields
     */
    public static boolean insert(String collectName, List<InsertParam.Field> fields, List<List<Float>> vectorList) {
        MilvusServiceClient milvusClient = MilvusConnectHelper.getMilvusServiceClient();
        log.info("milvus insert name:{}, Fields:{}, vectorList:{}", collectName, fields, vectorList);
        // 构建插入请求
        InsertParam.Builder insertParam = InsertParam.newBuilder()
                .withCollectionName(collectName)
                .withFields(fields);
        // 执行插入操作
        R<MutationResult> response = milvusClient.insert(insertParam.build());
        if (response.getStatus() != R.Status.Success.getCode()) {
            log.info("milvus insert vector fail! message:{}", response.getMessage());
            return false;
        }
        Exception exception = response.getException();
        if (exception != null) {
            log.error("从 milvus 中删除记录失败，cause {}", response.getMessage());
            return false;
        }
        MutationResult data = response.getData();
        // 已插入实体的条数
        data.getInsertCnt();
        List<Long> ids = data.getIDs().getIntId().getDataList();
        return true;
    }

    /**
     * 删除操作：根据ids删除向量
     *
     * @param collectName 指定集合名称
     * @param ids         包含要删除的实体们的主键
     */
    public static boolean delete(String collectName, List<Long> ids) {
        if (CollectionUtil.isEmpty(ids)) {
            return false;
        }
        String expr = "id in " + JSONObject.toJSONString(ids);
        MilvusServiceClient milvusClient = MilvusConnectHelper.getMilvusServiceClient();

        // 构建删除请求
        DeleteParam deleteParam = DeleteParam.newBuilder()
                .withCollectionName(collectName)
                .withExpr(expr)
                .build();
        // 执行删除操作
        R<MutationResult> response = milvusClient.delete(deleteParam);

        if (response == null) {
            log.error("从 milvus 中删除记录失败( ids = {} ) cause response is empty", JSONObject.toJSONString(ids));
            return false;
        }
        if (R.Status.Success.getCode() != response.getStatus()) {
            log.error("从 milvus 中删除记录( ids = {} )失败，失败原因是：{}", JSONObject.toJSONString(ids), response.getMessage());
            return true;
        }
        Exception exception = response.getException();
        if (exception != null) {
            log.error("从 milvus 中删除记录( ids = {} )异常，异常原因是：{}", JSONObject.toJSONString(ids), exception.getMessage());
            return false;
        }
        log.info("已删除实体的条数：{}", response.getData().getDeleteCnt());
        return false;
    }
}
